# python .\zip\zip_filetype.py e:\

import os
import zipfile
import tempfile
from pathlib import Path
import argparse

def extract_zip_recursively(zip_file_path, temp_dir):
    """递归地提取ZIP文件及其内部的ZIP文件到临时目录"""
    extracted_files = []

    def _extract_zip(_zip_file_path, _temp_dir):
        try:
            with zipfile.ZipFile(_zip_file_path, 'r') as zip_ref:
                zip_ref.extractall(_temp_dir)
                for zip_info in zip_ref.infolist():
                    if not zip_info.is_dir():
                        inner_file_path = Path(_temp_dir) / zip_info.filename
                        if inner_file_path.suffix.lower() == '.zip':
                            # 递归处理嵌套的ZIP文件
                            nested_temp_dir = Path(tempfile.mkdtemp(dir=_temp_dir))
                            _extract_zip(inner_file_path, nested_temp_dir)
                            extracted_files.extend(extract_zip_recursively(inner_file_path, nested_temp_dir))
                        else:
                            extracted_files.append((zip_file_path, zip_info.filename, inner_file_path))
        except zipfile.BadZipFile:
            print(f"警告: 无法打开ZIP文件 {_zip_file_path}")

    _extract_zip(zip_file_path, temp_dir)
    return extracted_files

def scan_directory(directory):
    """扫描目录，统计各类型文件的数量和占用空间"""
    ext_stats = {}

    def update_stats(file_path, file_name):
        try:
            # 获取文件大小
            file_size = os.path.getsize(file_path)
            # 获取文件扩展名
            _, ext = os.path.splitext(file_name)
            if not ext:  # 如果没有扩展名，默认为无扩展名''
                ext = 'no_extension'

            if ext in ext_stats:
                ext_stats[ext]['count'] += 1
                ext_stats[ext]['size'] += file_size
            else:
                ext_stats[ext] = {'count': 1, 'size': file_size}
        except Exception as e:
            print(f"无法处理文件 {file_path}: {e}")

    for root, _, files in os.walk(directory):
        for file_name in files:
            file_path = os.path.join(root, file_name)
            update_stats(file_path, file_name)

    return ext_stats

def format_bytes(size):
    """格式化字节大小为可读格式"""
    for unit in ['B', 'KB', 'MB', 'GB', 'TB']:
        if size < 1024:
            return f"{size:.2f} {unit}"
        size /= 1024
    return f"{size:.2f} PB"

def main(scan_dir):
    temp_dir = Path(tempfile.mkdtemp())
    all_files = []

    # 扫描指定目录下的所有ZIP文件
    for item in Path(scan_dir).rglob('*.zip'):
        if item.is_file():
            extracted_files = extract_zip_recursively(item, temp_dir)
            all_files.extend(extracted_files)

    # 统计解压出来的所有文件
    ext_stats = {}
    for zip_file_path, inner_file_name, actual_file_path in all_files:
        if actual_file_path.exists():  # 确保文件存在
            # 获取文件大小
            file_size = os.path.getsize(actual_file_path)
            # 获取文件扩展名
            _, ext = os.path.splitext(inner_file_name)
            if not ext:  # 如果没有扩展名，默认为无扩展名''
                ext = 'no_extension'

            if ext in ext_stats:
                ext_stats[ext]['count'] += 1
                ext_stats[ext]['size'] += file_size
            else:
                ext_stats[ext] = {'count': 1, 'size': file_size}

    # 输出结果
    print("文件类型统计：")
    for ext, info in sorted(ext_stats.items()):
        count = info['count']
        size = format_bytes(info['size'])
        print(f"{ext}: 数量={count}, 占用空间={size}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="扫描指定目录中的所有ZIP文件，统计其中的文件类型的数量和占用空间")
    parser.add_argument('--scan-dir', required=True, help="扫描目录")

    args = parser.parse_args()

    scan_dir = args.scan_dir

    main(scan_dir)